package com.zijidelu.luframework.utils.algo.es;

import cn.easyes.common.params.SFunction;
import cn.easyes.core.toolkit.FieldUtils;
import cn.hutool.core.util.ReflectUtil;
import com.zijidelu.luframework.utils.lang.function.Consumer3;
import com.zijidelu.luframework.utils.CastUtil;
import lombok.Data;
import lombok.RequiredArgsConstructor;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.search.aggregations.Aggregation;
import org.elasticsearch.search.aggregations.bucket.terms.ParsedTerms;
import org.elasticsearch.search.aggregations.bucket.terms.Terms;

import java.lang.reflect.Field;
import java.util.*;
import java.util.function.BiFunction;
import java.util.function.Predicate;
import java.util.function.UnaryOperator;
import java.util.stream.Collectors;

/**
 * Es聚合搜索算法。
 *
 * @author ZIJIDELU
 * @date: 2023/11/22 12:05
 */
@RequiredArgsConstructor
public class EsAggregationSearchAlgo<Record, Fill> {
    public static final String TERMS_SUFFIX = "Terms";
    private final Class<Record> recordClass;
    private final Class<Fill> fillClass;
    private boolean isDfs;
    private List<SFunction<Record, ?>> groupByColumnsList;
    private List<SFunction<Record, ?>> aggregationColumnsList;
    private SearchResponse searchResponse;
    private UnaryOperator<String> isLastLayer;
    private BiFunction<Terms.Bucket, Record, Fill> lastLayerOps;

    private static String suffixNameTerms(final String s) {
        return s + TERMS_SUFFIX;
    }

    private static <Record> List<String> toFieldNames(final List<SFunction<Record, ?>> sFunctions) {
        return sFunctions.stream()
                .map(FieldUtils::getFieldName)
                .collect(Collectors.toList());
    }

    /**
     * 获取所有字段Map。
     *
     * @param tClass 泛型类
     * @param <T>    泛型
     * @return 字段map
     */
    private static <T> Map<String, Field> getFieldMap(final Class<T> tClass) {
        return Arrays.stream(ReflectUtil.getFields(tClass))
                .peek(field -> field.setAccessible(true))
                .collect(Collectors.toMap(Field::getName, field -> field));
    }

    private static String getFieldName(final Terms terms) {
        final String name = terms.getName();
        return name.substring(0, name.indexOf(TERMS_SUFFIX));
    }

    /**
     * 递归算法，计算深层聚合数据。
     *
     * @param groupByTerms      分组EsTerms
     * @param groupByFieldNames 字段名称数组
     * @param i                 数组索引
     * @param fillObjectList    填充数据列表
     * @param recordObject      数据记录对象，只记录分组数据
     * @param isLastLayer       判断是否是最深层级
     * @param lastLayerConsumer 最深层级消费者，用于记录数据到列表
     * @return 记录数据列表
     */
    public static <Record, Fill> List<Fill>
    recursionAggregation(final ParsedTerms groupByTerms,
                         final String[] groupByFieldNames,
                         final int i,
                         final List<Fill> fillObjectList,
                         Record recordObject,
                         final Predicate<Terms.Bucket> isLastLayer,
                         final Consumer3<Terms.Bucket, Record, List<Fill>> lastLayerConsumer) {
        // 记录数据对象不允许为空
        assert Objects.nonNull(recordObject);
        final Class<?> recordObjectClass = recordObject.getClass();
        // 记录对象字段缓存
        final Map<String, Field> groupFieldMap = getFieldMap(recordObjectClass);
        // 基线条件
        // 判断索引是否越界，或这获取EsTerms为空。
        if (i >= groupByFieldNames.length || Objects.isNull(groupByTerms)) {
            return fillObjectList;
        }

        // 指针下移一位
        final int j = i + 1;

        for (Terms.Bucket bucket : groupByTerms.getBuckets()) {
            // 递归最外层创建对象记录分组数据
            if (i == 0) {
                recordObject = CastUtil.cast(ReflectUtil.newInstance(recordObjectClass));
            }

            // 分组字段写入值
            try {
                groupFieldMap.get(groupByFieldNames[i]).set(recordObject, bucket.getKeyAsString());
            } catch (IllegalAccessException e) {
                throw new RuntimeException(e);
            }

            // 判断是否进入递归最深层
            if (isLastLayer.test(bucket)) {
                // recordObject禁止写入fillObjectList
                // fillObjectList中的对象，应使用重新生成的对象且从recordObject中复制属性
                lastLayerConsumer.apply(bucket).apply(recordObject).accept(fillObjectList);
            }

            // 递归条件
            // 判断是否数组越界
            if (j < groupByFieldNames.length) {
                final ParsedTerms parsedTerms = bucket.getAggregations().get(suffixNameTerms(groupByFieldNames[j]));
                // 递归调用
                recursionAggregation(parsedTerms, groupByFieldNames, j, fillObjectList, recordObject, isLastLayer, lastLayerConsumer);
            }

        }

        return fillObjectList;
    }

    @Data
    private static class DfsBucket {
        private Terms.Bucket bucket;
        private String fieldName;
    }

    private static List<DfsBucket> listDfsBuckets(Terms terms) {
        return terms.getBuckets().stream()
                .map(bucket -> {
                    DfsBucket dfsBucket = new DfsBucket();
                    dfsBucket.setBucket(bucket);
                    dfsBucket.setFieldName(getFieldName(terms));
                    return dfsBucket;
                }).collect(Collectors.toList());
    }

    /**
     * 深度优先搜索算法，获取聚合数据。
     *
     * @param recordClass    记录对象类型
     * @param searchResponse es聚合结果
     * @param isLastLayer    判断是否到达最深层
     * @param lastLayerFunc  最深层函数转换
     * @param <Record>       记录对象泛型
     * @param <Fill>         填充对象泛型
     * @return 填充对象列表
     */
    public static <Record, Fill> List<Fill> dfsAggregation(final Class<Record> recordClass,
                                                           final SearchResponse searchResponse,
                                                           final Predicate<Terms.Bucket> isLastLayer,
                                                           final BiFunction<Terms.Bucket, Record, Fill> lastLayerFunc) {
        // 创建填充对象列表
        final List<Fill> fillList = new LinkedList<>();
        // 记录对象字段缓存
        final Map<String, Field> groupFieldMap = getFieldMap(recordClass);
        // 获取第一分组terms
        final List<Aggregation> aggregationList = searchResponse.getAggregations().asList();
        for (Aggregation aggregation : aggregationList) {
            final Terms terms = ((Terms) aggregation);
            for (DfsBucket dfsBucket : listDfsBuckets(terms)) {
                // 使用双端队列，模拟栈。
                final Deque<DfsBucket> bucketDeque = new LinkedList<>();
                // 第一分组bucket入栈。
                bucketDeque.push(dfsBucket);
                // 创建记录对象
                final Record record = ReflectUtil.newInstance(recordClass);
                while (!bucketDeque.isEmpty()) {
                    // 弹栈
                    final DfsBucket popDfsBucket = bucketDeque.pop();
                    // 设置记录对象字段的值
                    final String groupValue = popDfsBucket.getBucket().getKeyAsString();
                    try {
                        groupFieldMap.get(popDfsBucket.getFieldName()).set(record, groupValue);
                    } catch (IllegalAccessException e) {
                        throw new RuntimeException(e);
                    }

                    // 判断最深层
                    final boolean isLastLayerValue = isLastLayer.test(popDfsBucket.getBucket());
                    if (isLastLayerValue) {
                        // 最深层操作
                        fillList.add(lastLayerFunc.apply(popDfsBucket.getBucket(), record));
                    }

                    // 获取子分组
                    final List<Aggregation> sonAggregationList = popDfsBucket.getBucket().getAggregations().asList();
                    for (Aggregation sonAggregation : sonAggregationList) {
                        // 不是最深层则允许展开
                        if (!isLastLayerValue) {
                            final Terms sonTerms = ((Terms) sonAggregation);
                            for (DfsBucket sonTermsDfsBucket : listDfsBuckets(sonTerms)) {
                                bucketDeque.push(sonTermsDfsBucket);
                            }
                        }
                    }

                }
            }
        }

        return fillList;
    }

    /**
     * 获取Es聚合数据，递归。
     * <p>
     * fill对象应使用new方式生成，或者MapStruct将record对象转换为fill对象。
     *
     * @param searchResponse         Es返回数据
     * @param groupByColumnsList     分组字段列表
     * @param aggregationColumnsList 聚合字段列表
     * @param recordObj              记录对象
     * @param isLastLayer            最深层聚合字段Es名称转换，验证作用
     * @param lastLayerFunc          转换填充对象Func
     * @param <Record>               记录对象泛型
     * @param <Fill>                 填充对象泛型
     * @return 聚合数据
     */
    public static <Record, Fill> List<Fill>
    listEsAggregationByRecurse(final SearchResponse searchResponse,
                               final List<String> groupByColumnsList,
                               final List<String> aggregationColumnsList,
                               final Record recordObj,
                               final UnaryOperator<String> isLastLayer,
                               final BiFunction<Terms.Bucket, Record, Fill> lastLayerFunc) {
        // 获取字段名称数组
        final String[] groupByFieldNames = groupByColumnsList.toArray(new String[0]);

        // 获取最外层分组
        final ParsedTerms parsedTerms = searchResponse.getAggregations().get(suffixNameTerms(groupByFieldNames[0]));

        return recursionAggregation(parsedTerms, groupByFieldNames, 0, new LinkedList<>(), recordObj,
                bucket -> isLastLayer(aggregationColumnsList, bucket, isLastLayer),
                bucket -> recordObject -> fillObjectList -> fillObjectList.add(lastLayerFunc.apply(bucket, recordObject))
        );
    }

    /**
     * 获取Es聚合数据，深度优先搜索。
     * <p>
     * fill对象应使用new方式生成，或者MapStruct将record对象转换为fill对象。
     *
     * @param searchResponse         Es返回数据
     * @param aggregationColumnsList 聚合字段列表
     * @param recordClass            记录对象类型
     * @param isLastLayer            最深层聚合字段Es名称转换，验证作用
     * @param lastLayerFunc          转换填充对象Func
     * @param <Record>               记录对象泛型
     * @param <Fill>                 填充对象泛型
     * @return 聚合数据
     */
    public static <Record, Fill> List<Fill>
    listEsAggregationByDfs(final SearchResponse searchResponse,
                           final List<String> aggregationColumnsList,
                           final Class<Record> recordClass,
                           final UnaryOperator<String> isLastLayer,
                           final BiFunction<Terms.Bucket, Record, Fill> lastLayerFunc) {
        return dfsAggregation(recordClass,
                searchResponse,
                bucket -> isLastLayer(aggregationColumnsList, bucket, isLastLayer),
                lastLayerFunc
        );
    }

    /**
     * 判断Bucket内部是否存在聚合数据，若存在则表明进入最深层Bucket。
     *
     * @param aggregationColumnsList            聚合字段列表
     * @param bucket                            es桶
     * @param fieldNameAggregationUnaryOperator 最深层聚合字段Es名称转换，验证作用
     */
    public static boolean isLastLayer(final List<String> aggregationColumnsList,
                                      final Terms.Bucket bucket,
                                      final UnaryOperator<String> fieldNameAggregationUnaryOperator) {
        for (String aggregationColumn : aggregationColumnsList) {
            if (Objects.nonNull(bucket.getAggregations().get(fieldNameAggregationUnaryOperator.apply(aggregationColumn)))) {
                return true;
            }
        }
        return false;
    }

    /**
     * 获取Es聚合数据。
     * <p>
     * fill对象应使用new方式生成，或者MapStruct将record对象转换为fill对象。
     *
     * @param searchResponse                    Es返回数据
     * @param groupByColumnsList                分组字段列表
     * @param aggregationColumnsList            聚合字段列表
     * @param recordObj                         记录对象
     * @param fieldNameAggregationUnaryOperator 最深层聚合字段Es名称转换，验证作用
     * @param fieldFunc                         转换填充对象Func
     * @param <Record>                          记录对象泛型
     * @param <Fill>                            填充对象泛型
     * @return 聚合数据
     */
    public static <Record, Fill> List<Fill>
    listEsAggregationLambdaField(final SearchResponse searchResponse,
                                 final List<SFunction<Record, ?>> groupByColumnsList,
                                 final List<SFunction<Record, ?>> aggregationColumnsList,
                                 final Record recordObj,
                                 final UnaryOperator<String> fieldNameAggregationUnaryOperator,
                                 final BiFunction<Terms.Bucket, Record, Fill> fieldFunc) {
        return listEsAggregationByRecurse(
                searchResponse,
                toFieldNames(groupByColumnsList),
                toFieldNames(aggregationColumnsList),
                recordObj,
                fieldNameAggregationUnaryOperator,
                fieldFunc
        );

    }

    public EsAggregationSearchAlgo<Record, Fill> dfs() {
        this.isDfs = true;
        return this;
    }

    public EsAggregationSearchAlgo<Record, Fill> recurse() {
        this.isDfs = false;
        return this;
    }

    public EsAggregationSearchAlgo<Record, Fill> groupColumns(final List<SFunction<Record, ?>> groupByColumnsList) {
        this.groupByColumnsList = groupByColumnsList;
        return this;
    }

    public EsAggregationSearchAlgo<Record, Fill> aggregationColumns(final List<SFunction<Record, ?>> aggregationColumnsList) {
        this.aggregationColumnsList = aggregationColumnsList;
        return this;
    }

    public EsAggregationSearchAlgo<Record, Fill> searchResponse(final SearchResponse searchResponse) {
        this.searchResponse = searchResponse;
        return this;
    }

    public EsAggregationSearchAlgo<Record, Fill> isAggregation(final UnaryOperator<String> isLastLayer) {
        this.isLastLayer = isLastLayer;
        return this;
    }

    public EsAggregationSearchAlgo<Record, Fill> lastLayerOps(final BiFunction<Terms.Bucket, Record, Fill> lastLayerOps) {
        this.lastLayerOps = lastLayerOps;
        return this;
    }


    public List<Fill> calculate() {
        final List<String> groupFieldNames = groupByColumnsList.stream()
                .map(FieldUtils::getFieldName)
                .collect(Collectors.toList());
        final List<String> aggregationFieldNames = aggregationColumnsList.stream()
                .map(FieldUtils::getFieldName)
                .collect(Collectors.toList());

        // dfs算法
        if (isDfs) {
            return listEsAggregationByDfs(searchResponse, aggregationFieldNames, recordClass, isLastLayer, lastLayerOps);
        }

        return listEsAggregationByRecurse(searchResponse, groupFieldNames, aggregationFieldNames,
                ReflectUtil.newInstance(recordClass), isLastLayer, lastLayerOps);
    }

}
